﻿using System.Collections.Generic;
using System.Runtime.InteropServices;
using Unity.Burst;
using Unity.Collections;
using Unity.Collections.LowLevel.Unsafe;
using Unity.Jobs;
using Unity.Mathematics;
using UnityEngine;
using UnityEngine.Rendering;
using Random = Unity.Mathematics.Random;

namespace Learning
{
    /// <summary>
    /// 使用 Job 和指针，优化 Simple2GPUInstancing 的代码
    /// </summary>
    public class SimpleUnsafe : MonoBehaviour
    {
        /// <summary>
        /// GPU Instancing 一个批次可提交的数量
        /// 注意：这里设置得越大，DC会越少，但是特殊情况下浪费的带宽会越多（比如只提交10个时，一个批次还是会提交512个数据给GPU，但是有效数据只有10个）
        /// </summary>
        private const int BATCH = 512;
        
        private const float Range = 100000f;
        private readonly Bounds BOUNDS = new Bounds(Vector3.zero, new Vector3(Range, Range, Range));
        
        public int instanceCount = 100000;
        public Mesh instanceMesh;
        public Material instanceMaterial;
        
        private NativeArray<float4x4> _transArray; // 每个对象对应的矩阵
        private NativeArray<float3> _forwardArray; // 每个对象1秒的前进增量
        private NativeArray<float3> _rotateArray; // 每个对象1秒的旋转增量
        
        private InitJob _initJob;
        private UpdateJob _updateJob;
        private JobHandle _updateHandle;

        private List<UnsafeBatchInfo> _pool = new List<UnsafeBatchInfo>();
        
        void Start()
        {
            _transArray = new NativeArray<float4x4>(instanceCount, Allocator.Persistent);
            _forwardArray = new NativeArray<float3>(instanceCount, Allocator.Persistent);
            _rotateArray = new NativeArray<float3>(instanceCount, Allocator.Persistent);

            _initJob = new InitJob();
            _initJob.transArray = _transArray;
            _initJob.forwardArray = _forwardArray;
            _initJob.rotateArray = _rotateArray;
            
            _updateJob = new UpdateJob();
            _updateJob.transArray = _transArray;
            _updateJob.forwardArray = _forwardArray;
            _updateJob.rotateArray = _rotateArray;
            
            var handle = _initJob.Schedule(instanceCount, 64);
            handle.Complete();
        }

        void Update()
        {
            _updateJob.time = Time.deltaTime;
            _updateHandle = _updateJob.Schedule(instanceCount, 64);
            JobHandle.ScheduleBatchedJobs();
        }

        private unsafe void LateUpdate()
        {
            _updateHandle.Complete();
            
            var batchCount = (int) math.ceil((float) instanceCount / (float) BATCH);
            for (int i = 0; i < batchCount; i++)
            {
                var startIndex = i * BATCH;
                var len = BATCH;
                if (i == batchCount - 1)
                {
                    len = instanceCount - (BATCH * i);
                }

                if (_pool.Count == i)
                {
                    var newInfo = new UnsafeBatchInfo();
                    
                    newInfo.batchTransArray = new NativeArray<float4x4>(BATCH, Allocator.Persistent);
                    newInfo.transBuf = new GraphicsBuffer(GraphicsBuffer.Target.Structured, BATCH, Marshal.SizeOf(typeof(float4x4)));
                    newInfo.transBuf.SetData(newInfo.batchTransArray);
                    
                    newInfo.block = new MaterialPropertyBlock();
                    newInfo.block.SetBuffer("transBuf", newInfo.transBuf);
                    
                    newInfo.args = new uint[5] {0, 0, 0, 0, 0};
                    newInfo.args[0] = instanceMesh.GetIndexCount(0); // 每个实例的索引数
                    newInfo.args[1] = 0; // 实例数
                    newInfo.args[2] = instanceMesh.GetIndexStart(0); // 起始索引位置
                    newInfo.args[3] = instanceMesh.GetBaseVertex(0); // 基顶点位置
                    newInfo.args[4] = 0; // 起始实例位置
                    newInfo.argsBuf = new GraphicsBuffer(GraphicsBuffer.Target.IndirectArguments, 1, 5 * sizeof(uint));
                    newInfo.argsBuf.SetData(newInfo.args);
                    
                    _pool.Add(newInfo);
                }
                
                var info = _pool[i];
                
                // 填充数据
                var sourcePtr = (float4x4*) _transArray.GetUnsafePtr();
                sourcePtr += startIndex;
                var destPtr = (float4x4*) info.batchTransArray.GetUnsafePtr();
                UnsafeUtility.MemCpy(destPtr, sourcePtr, BATCH * 16 * 4); // 直接通过指针拷贝数据
                info.transBuf.SetData(info.batchTransArray); // 数据变动后需要重新设置才能生效
                
                info.args[1] = (uint) len; // 实例数
                info.argsBuf.SetData(info.args); // 数据变动后需要重新设置才能生效
                
                // 绘制
                Graphics.DrawMeshInstancedIndirect(instanceMesh, 0, instanceMaterial, BOUNDS, info.argsBuf,
                    0, info.block, ShadowCastingMode.On, true, LayerMask.NameToLayer("Character"));
            }
        }
    }

    public class UnsafeBatchInfo
    {
        public MaterialPropertyBlock block;
        public GraphicsBuffer transBuf;
        public NativeArray<float4x4> batchTransArray;
        public GraphicsBuffer argsBuf;
        public uint[] args;
    }
    
    /// <summary>
    /// 初始化所有对象，安全代码
    /// </summary>
    [BurstCompile]
    public struct InitJob : IJobParallelFor
    {
        public NativeArray<float4x4> transArray;
        public NativeArray<float3> forwardArray;
        public NativeArray<float3> rotateArray;
        
        public void Execute(int i)
        {
            var random = Random.CreateFromIndex((uint) i);
            transArray[i] = float4x4.TRS(
                    random.NextFloat3(-1000, 1000),
                    quaternion.Euler(random.NextFloat3(0, 360)),
                    random.NextFloat3(1, 5)
                );
            forwardArray[i] = random.NextFloat3(0, 1);
            rotateArray[i] = random.NextFloat3(0, 10);
        }
    }
    
    /// <summary>
    /// 每帧运算新的坐标旋转信息，非安全代码
    /// </summary>
    [BurstCompile]
    public struct UpdateJob : IJobParallelFor
    {
        public NativeArray<float4x4> transArray;
        // 标记为 ReadOnly 后就不能使用指针了 `forwardArray.GetUnsafePtr()` 这句代码会报错，会被认为可以修改其内部数据
        // [ReadOnly]
        public NativeArray<float3> forwardArray;
        // [ReadOnly]
        public NativeArray<float3> rotateArray;
        [ReadOnly]
        public float time;
        
        public unsafe void Execute(int i)
        {
            var transArrPtr = (float4x4*) transArray.GetUnsafePtr(); // 不需要 fixed 语句，因为 NativeArray 就是非托管的资源
            var transPtr = transArrPtr + i; // 通过指针直接取数组中的元素，效率比 arr[i] 更快

            var forwardArrPtr = (float3*) forwardArray.GetUnsafePtr();
            var forwardPtr = forwardArrPtr + i;

            var rotateArrPtr = (float3*) rotateArray.GetUnsafePtr();
            var rotatePtr = rotateArrPtr + i;

            // 如果只有位移，可以用下面的代码，更简单
            // transPtr->c3.xyz += *forwardPtr * time;

            // 存在位移和旋转时
            var rotate = new quaternion(*transPtr);
            rotate = math.mul(rotate, quaternion.Euler(*rotatePtr * time));
            var mat = new float4x4(rotate, *forwardPtr * time);
            // *transPtr *= mat; // 注意 Mathematics 里没有重置操作符，所以不要像 Matrix4x4 那样直接使用 * 号运算，而是要调用 math.mul 方法
            *transPtr = math.mul(*transPtr, mat);
        }
    }
}
